import cv2
import numpy as np
import time
import math

trace_stack = []
side_stack = [0,0,0,0,0,0]

def kn(n):
    return np.ones((n,n)).astype(np.uint8)

def blue_mask(hsv):
    low_thres = np.array([150,60,30])
    high_thres = np.array([230,255,230])
    mask =  cv2.inRange(hsv,low_thres,high_thres)
    return cv2.morphologyEx(mask, cv2.MORPH_OPEN, kn(5))

def red_mask(hsv):
    low_thres = np.array([240,160,60])
    high_thres = np.array([255,255,240])        
    mask1 = cv2.inRange(hsv,low_thres,high_thres)
    low_thres = np.array([0,160,60])
    high_thres = np.array([30,255,240])        
    mask2 = cv2.inRange(hsv,low_thres,high_thres)
    mask =  cv2.bitwise_or(mask1,mask2)
    return cv2.morphologyEx(mask, cv2.MORPH_OPEN, kn(5))

def vec_calc_len(vec):
    x = vec[0]
    y = vec[1]
    return math.sqrt(x*x + y*y)

def vecs_calc_included_cos(vec1,vec2):
    dot = vec1[0]*vec2[0] + vec1[1]*vec2[1]
    lv1 = vec_calc_len(vec1)
    lv2 = vec_calc_len(vec2)
    min_len = 12
    if if_inRange(lv1,0,min_len) or if_inRange(lv2,0,min_len):
        return (2,lv2)
    else:
        return (dot/(lv1 * lv2),lv2)

def if_inRange(value,b1,b2):
    if b1 < b2:
        bottom = b1
        top = b2
    else:
        bottom = b2
        top = b1
    if value >= bottom and value <= top:
        return True
    else:
        return False
    
def if_near_border(point,w,h,ref = 10):
    x = point[0]
    y = point[1]
    res = if_inRange(x,0,ref) or if_inRange(x,w,w-ref) or if_inRange(y,0,ref) or if_inRange(y,h,h-ref)
    return res

def corner_analyz(polycor,w,h):
    length = len(polycor)
    cos_list = []
    d_list = []
    attach_border = False
    if length < 7:
        for i in range(length):
            s = i+1
            if s >= length:
                s -= length
            point_off_area = if_near_border(polycor[i][0],w,h)
            if not point_off_area:
                v = [polycor[i][0][0]-polycor[s][0][0],polycor[i][0][1]-polycor[s][0][1]]
                vp = [polycor[i][0][0]-polycor[i-1][0][0],polycor[i][0][1]-polycor[i-1][0][1]]
                v_cos,d = vecs_calc_included_cos(v,vp)
                if not v_cos == 2:
                    cos_list.append(v_cos)
                    d_list.append(d)
            else:
                attach_border = True
        mean_cos = np.mean(np.array(cos_list))
        var = np.var(np.array(d_list))
        # print(h/length/4)
        if var < h/length/4 or attach_border:
            if if_inRange(mean_cos,0.64278,0.34202): # cos50~cos70
                return 3
            elif if_inRange(mean_cos,0.17365,-0.17365): # cos80~cos100
                return 4
            elif if_inRange(mean_cos,-0.5,-0.9848): # cos120~cos170
                return 0
            else:
                #time.sleep(0.2)
                #print('D|len:{},cos:{},var:{}'.format(length,mean_cos,var))
                return 1
        else:
            # time.sleep(0.2)
            # print('O|len:{},cos:{},var:{}'.format(length,mean_cos,var))
            # print(d_list)
            return 1
    elif length > 15:
        return 1
    else:
        return 0

def draw_shape(mask,img,sp):
    contours,_ = cv2.findContours(mask,cv2.RETR_LIST,cv2.CHAIN_APPROX_NONE)
    h,w = mask.shape
    for ct in contours:
        if len(ct) > 150:
            peri = cv2.arcLength(ct,True)
            approx = cv2.approxPolyDP(ct, 0.02*peri, True)
            x = approx[:,:,0].flatten()
            y = approx[:,:,1].flatten()
            ox = np.mean(x)+1
            oy = np.mean(y)+1

            side = corner_analyz(approx,w,h)
            font = cv2.FONT_HERSHEY_SIMPLEX
            pos = (int(ox)-20,int(oy)+10)

            cv2.putText(img,'{}{}'.format(sp,side),pos,font,1,(0,255,255),2,8)
            cv2.drawContours(img,ct,-1,(0,255,0),2)
            for corner in approx:
                cx = corner[0][0]
                cy = corner[0][1]
                cv2.circle(img,(cx,cy),4,(255,20,20),3)

def find_shape(mask,img,cors):
    contours,_ = cv2.findContours(mask,cv2.RETR_LIST,cv2.CHAIN_APPROX_NONE)
    h,w = mask.shape
    tar_poss = []
    max_ct = []

    for ct in contours:
        if len(ct) > 150:
            peri = cv2.arcLength(ct,True)
            approx = cv2.approxPolyDP(ct, 0.02*peri, True)
            x = approx[:,:,0].flatten()
            y = approx[:,:,1].flatten()
            ox = np.mean(x)+1
            oy = np.mean(y)+1

            side = corner_analyz(approx,w,h)
            if side == cors:
                #cv2.circle(img,(int(ox),int(oy)),5,(0,255,255),2)
                cv2.drawContours(img,ct,-1,(255,0,255),5)
                tar_poss.append([ox,oy])
            
            elif cors == -1:
                len_max_ct = len(max_ct)
                if len_max_ct == 0:
                    max_ct = ct
                elif len_max_ct < len(ct):
                    max_ct = ct

    if cors == -1:
        if len(max_ct) == 0:
            return 1
        
        else:
            peri = cv2.arcLength(max_ct,True)
            approx = cv2.approxPolyDP(max_ct, 0.02*peri, True)
            return corner_analyz(approx,w,h)

    else:      
        return tar_poss

def update_stack(pos):
    global trace_stack
    length = len(trace_stack)
    max_len = 20
    #print(trace_stack)
    if length >= max_len: 
        trace_stack = trace_stack[1:]

    if length == 0:
        trace_stack.append(pos)
    else:
        end = trace_stack[-1]
        vec = [pos[0]-end[0],pos[1]-end[1]]
        d = vec_calc_len(vec)
        if d < 20:
            trace_stack.append(pos)
        else:
            trace_stack = trace_stack[:-1]
    length = len(trace_stack)
    if length > max_len/2:
        return (True,pos)
    else:
        return (False,None)

def draw_path(img,path):
    length = len(path)
    step = int(253/length)
    factor = range(0,step*length,step)
    for i in range(length):
        x = int(path[i][0])
        y = int(path[i][1])
        cv2.circle(img,(x,y),1,(255-factor[i],factor[i],255),2)

def clear_stack():
    global trace_stack
    trace_stack = []

def decide_target(img):
    global side_stack
    
    img = cv2.GaussianBlur(img,(3,3),4)
    hsv = cv2.cvtColor(img,cv2.COLOR_BGR2HSV_FULL )
    mask_b = blue_mask(hsv)
    mask_r = red_mask(hsv)
    side_b = find_shape(mask_b,img,-1)
    side_r = find_shape(mask_r,img,-1)

    m = [0,0,0,1,2]
    if not side_b == 1:
        side_stack[m[side_b]] += 1

    if not side_r == 1:
        side_stack[m[side_r] + 3] += 1

    im = ['b0','b3','b4','r0','r3','r4']
    # print(side_stack)
    max_stack = max(side_stack)
    for i in range(len(side_stack)):
        if side_stack[i] == max_stack and max_stack > 20:
            return im[i]

    return 'None' 


def find_target(img,tar):
    global trace_stack
    img = cv2.GaussianBlur(img,(3,3),4)
    hsv = cv2.cvtColor(img,cv2.COLOR_BGR2HSV_FULL )
    tar_color = tar[:1]
    tar_sides = int(tar[1:])
    if tar_color == 'b':
        mask = blue_mask(hsv)
    elif tar_color == 'r':
        mask = red_mask(hsv)
    else:
        mask = None
    poss = find_shape(mask,img,tar_sides)

    draw_shape(mask,img,tar_color)
    #---------------------
    # mask_red = red_mask(hsv)
    # mask_blue = blue_mask(hsv)
    # draw_shape(mask_red,img,'r')
    # draw_shape(mask_blue,img,'b')
    #-----------------------
    # print(len(poss))
    for pos in poss:
        ret,respos = update_stack(pos)
        if ret:
            draw_path(img,trace_stack)
            return (img,respos)

    return (img,[-1,-1])

def dot_mask(hsv):
    low_thres = np.array([230,60,180])
    high_thres = np.array([255,255,255]) 
    red1_mask = cv2.inRange(hsv,low_thres,high_thres)
    low_thres = np.array([0,60,180])
    high_thres = np.array([50,255,255]) 
    red2_mask = cv2.inRange(hsv,low_thres,high_thres)
    red_mask = cv2.bitwise_or(red1_mask,red2_mask)
    mask = cv2.morphologyEx(red_mask, cv2.MORPH_CLOSE, kn(7))
    low_thres = np.array([0,50,160])
    high_thres = np.array([255,255,245]) 
    light_mask = cv2.inRange(hsv,low_thres,high_thres)
    mask = cv2.bitwise_xor(mask,light_mask)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kn(5))
    return mask

def find_dot(mask,img):
    contours,_ = cv2.findContours(mask,cv2.RETR_LIST,cv2.CHAIN_APPROX_NONE)
    got = []
    h,w,_ = img.shape
    for ct in contours:
        if (len(ct)>5 and len(ct)<60): 
            x = ct[:,:,0].flatten()
            y = ct[:,:,1].flatten()
            if np.var(x)>5 and np.var(y)>5:
                dx = np.max(x) - np.min(x)
                dy = np.max(y) - np.min(y)
                ox = np.mean(x)+1
                oy = np.mean(y)+1
                r = (dx+dy)/4
                #print([ox,oy],r)
                if r >= 3 and abs(dx-dy) < 5 and abs(ox-w/2-20)<20 and abs(oy-h/2-20)<25:
                    cv2.circle(img,(int(ox),int(oy)),int(r),(255,0,0),2)
                    got.append(ct)
                    img = cv2.drawContours(img,ct,-1,(0,255,255),1)

#=======================================================================

# 0是代表摄像头编号，只有一个的话默认为0

# capture = cv2.VideoCapture('./video/shape2.mp4')
# capture.set(cv2.CAP_PROP_FPS, 15)

# print('start')
# now = time.time()
# init = 0
# fourcc=0x7634706d

# tar = 'r3'

# while True:
#     # 调用摄像机
#     ref, img = capture.read()
#     #img = np.array(img).astype(np.float32)
#     last = now
#     now = time.time()
#     if now!=last:
#         fps = 1/(now - last)
#         # print('Fps:{}'.format(fps))
#     else: 
#         fps = 0

#     if (ref):
        
#         res = decide_target(img)
#         img,pos = find_target(img,tar)

#         print(res)


#         if(init == 0):
#             #writer = cv2.VideoWriter('./video/detect_res_{}.mp4'.format(time.localtime().tm_min),fourcc,15,(img.shape[1], img.shape[0]))
#             init = 1
#         elif(init == 1): 
#             #writer.write(img)
#             pass
#         init = 1

#         h,w,_ = img.shape
#         img = cv2.resize(img,(int(w/2),int(h/2)))
#         cv2.imshow('PC Camera', img)

#     c = cv2.waitKey(5) & 0xff
#     if (ref == False and init == 1 or c == 27):
#         # 释放所有窗口
#         cv2.destroyAllWindows()
#         break